#pragma once
#include <iostream>
#include <cassert>
using namespace std;

template<class K, class V>
struct avl_node
{
    avl_node<K, V>* _left;
    avl_node<K, V>* _right;
    avl_node<K, V>* _parent;

    pair<K, V> _kv;
    int _bf; // balance factor

    avl_node<K, V>(const pair<K,V>& kv)
        : _left(nullptr), _right(nullptr), _parent(nullptr)
        , _kv(kv), _bf(0)
    {}
};

template<class K, class V>
class avl_tree
{
public:
    typedef avl_node<K, V> node;

public:
    void inorder()
    {
        _inorder(_root);
        cout << endl;
    }
    void _inorder(node* root)
    {
        if (!root) return;

        _inorder(root->_left);
        cout << root->_kv.first << ":" << root->_kv.second << " ";
        _inorder(root->_right);
    }

    bool is_balance()
    {
        return _is_balance(_root);
    }
    bool _is_balance(node* root)
    {
        if (!root)
            return true;

        int lh  = height(root->_left);
        int rh = height(root->_right);

        if (rh - lh != root->_bf)
        {
            cout << root->_kv.first << "  but now:" << root->_bf << endl;
            cout << root->_kv.first << "should be:" << rh - lh << endl;
            return false;
        }

        return abs(rh - lh) < 2
            && _is_balance(root->_left)
            && _is_balance(root->_right);
    }

    int height(node* root)
    {
        if (!root)
            return 0;

        int lh = height(root->_left);
        int rh = height(root->_right);

        return rh > lh ? rh + 1 : lh + 1;
    }

public:
    bool insert(const pair<K, V>& kv)
    {
        if (!_root)
        {
            _root = new node(kv);
            return true;
        }

        node* parent = nullptr;
        node* cur = _root;

        while (cur)
        {
            if (cur->_kv.first < kv.first)
            {
                parent = cur;
                cur = cur->_right;
            }
            else if (cur->_kv.first > kv.first)
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return false;
            }
        }

        cur = new node(kv);
        if (parent->_kv.first < kv.first)
            parent->_right = cur;
        else
            parent->_left = cur;
        cur->_parent = parent; // maintain parent ptr

        // update balance factor
        while (parent)
        {
            if (cur == parent->_left)
                --parent->_bf;
            else
                ++parent->_bf;

            if (parent->_bf == 0)
            {
                break;
            }
            else if (parent->_bf == 1 || parent->_bf == -1)
            {
                cur = parent;
                parent = cur->_parent;
            }
            else if (parent->_bf == 2 || parent->_bf == -2) // rotate
            {
                if (cur->_bf == 1 && parent->_bf == 2)
                    rotate_l(parent);
                else if (cur->_bf == -1 && parent->_bf == -2)
                    rotate_r(parent);
                else if (cur->_bf == 1 && parent->_bf == -2)
                    rotate_lr(parent);
                else if (cur->_bf == -1 && parent->_bf == 2)
                    rotate_rl(parent);
                else
                    assert(false);
                break;
            }
            else
            {
                assert(false);
            }
        }

        return true;
    }

private:
    void rotate_l(node* parent)
    {
        node* subr = parent->_right;
        node* subrl = subr->_left;

        parent->_right = subrl;
        if (subrl) subrl->_parent = parent;

        node* pparent = parent->_parent;

        subr->_left = parent;
        parent->_parent = subr;

        if (parent == _root)
            _root = subr;
        else
        {
            if (pparent->_left == parent)
                pparent->_left = subr;
            else
                pparent->_right = subr;
        }
        subr->_parent = pparent;

        parent->_bf = 0;
        subr->_bf = 0;
    }

    void rotate_r(node* parent)
    {
        node* subl = parent->_left;
        node* sublr = subl->_right;

        parent->_left = sublr;
        if (sublr) sublr->_parent = parent;

        node* pparent = parent->_parent;

        subl->_right = parent;
        parent->_parent = subl;

        if (parent == _root)
        {
            _root = subl;
            subl->_parent = nullptr;
        }
        else
        {
            if (pparent->_left == parent)
                pparent->_left = subl;
            else
                pparent->_right = subl;
            subl->_parent = pparent;
        }

        parent->_bf = 0;
        subl->_bf = 0;
    }

    void rotate_lr(node* parent)
    {
        node* subl = parent->_left;
        node* sublr = subl->_right;

        int bf = sublr->_bf;

        rotate_l(parent->_left);
        rotate_r(parent);

        if (bf == 0)
        {
            parent->_bf = 0;
            subl->_bf = 0;
            sublr->_bf = 0;
        }
        else if (bf == 1)
        {
            parent->_bf = 0;
            subl->_bf = -1;
            sublr->_bf = 0;
        }
        else if (bf == -1)
        {
            parent->_bf = 1;
            subl->_bf = 0;
            sublr->_bf = 0;
        }
        else
        {
            assert(false);
        }
    }

    void rotate_rl(node* parent)
    {
        node* subr = parent->_right;
        node* subrl = subr->_left;

        int bf = subrl->_bf;

        rotate_r(parent->_right);
        rotate_l(parent);

        if (bf == 0)
        {
            parent->_bf = 0;
            subr->_bf = 0;
            subrl->_bf = 0;
        }
        else if (bf == -1)
        {
            parent->_bf = 0;
            subr->_bf = 1;
            subrl->_bf = 0;
        }
        else if (bf == 1)
        {
            parent->_bf = -1;
            subr->_bf = 0;
            subrl->_bf = 0;
        }
        else
        {
            assert(false);
        }
    }

private:
    node* _root = nullptr;
};